package middleware

import (
	"mcode/pkg/response"
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
)

// 简单的限流器实现
type rateLimiter struct {
	requests map[string][]time.Time
	mu       sync.RWMutex
	max      int
	window   time.Duration
}

func newRateLimiter(max int, window time.Duration) *rateLimiter {
	return &rateLimiter{
		requests: make(map[string][]time.Time),
		max:      max,
		window:   window,
	}
}

func (rl *rateLimiter) allow(key string) bool {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	now := time.Now()
	cutoff := now.Add(-rl.window)

	// 清理过期的请求
	if times, exists := rl.requests[key]; exists {
		var validTimes []time.Time
		for _, t := range times {
			if t.After(cutoff) {
				validTimes = append(validTimes, t)
			}
		}
		rl.requests[key] = validTimes
	}

	// 检查是否超过限制
	if len(rl.requests[key]) >= rl.max {
		return false
	}

	// 记录新请求
	rl.requests[key] = append(rl.requests[key], now)
	return true
}

// RateLimit 返回一个限流中间件
// max: 在指定时间窗口内的最大请求数
// expiration: 时间窗口大小
func RateLimit(max int, expiration time.Duration) gin.HandlerFunc {
	limiter := newRateLimiter(max, expiration)

	return func(c *gin.Context) {
		key := c.ClientIP() // 使用IP作为限流的key

		if !limiter.allow(key) {
			response.Error(c, http.StatusTooManyRequests, "请求太频繁，请稍后再试")
			c.Abort()
			return
		}

		c.Next()
	}
}
